"""
@Time       : 2025/9/22 14:31
@Author     : cholise26
@File       : doubao_generate_iamge.py
@Software   : PyCharm 
"""
import requests
import os
from typing import List, Dict
from datetime import datetime
import json

# -------------------------- 1. 核心参数配置 --------------------------
API_KEY = "sk-PjfrbJ5yZ1kAeSH6875c4fD5551f4f4b8e839d28Da91A6Db"  # 替换为你的豆包正确密钥
MODEL_ID = "doubao-seedream-4-0-250828"

# 待生成的图片描述列表（每条描述将与最终地址关联）
IMAGE_PROMPTS = [
    "A close-up of a delicate gold chain necklace lying on a soft, neutral fabric surface, showcasing its clean lines and subtle shine, with natural light illuminating its texture.",
]

# 豆包API接口参数
API_URL = f"https://aihubmix.com/v1/models/doubao/{MODEL_ID}/predictions"
GEN_PARAMS = {
    "size": "2K",
    "sequential_image_generation": "disabled",
    "stream": False,
    "response_format": "url",
    "watermark": True
}


# -------------------------- 2. 工具函数（已不再需要Base64解码） --------------------------
# 原有的 base64_to_image 函数被移除，因为豆包API直接返回URL。
# init_output_dir 函数也不再需要，因为不进行本地保存。

# -------------------------- 3. 单图生成函数（返回图片URL） --------------------------
def generate_single_image(prompt: str) -> str or None:
    """成功返回图片URL，失败返回None"""
    try:
        print(f"\n📌 处理描述：{prompt[:50]}...")
        prompt = "严格按描述生成16:9的高清图片，不要出现水印：" + prompt

        # 构造请求头
        headers = {
            "Content-Type": "application/json",
            "Authorization": f"Bearer {API_KEY}"
        }

        # 构造请求体
        payload = {
            "input": {
                "prompt": prompt,
                **GEN_PARAMS
            }
        }

        # 发送POST请求
        response = requests.post(API_URL, headers=headers, json=payload, timeout=60)

        # 检查响应状态码
        response.raise_for_status()

        # 解析返回的JSON
        result = response.json()

        # 提取图片URL
        if "output" in result and result["output"]:
            image_url = result["output"][0].get("url")
            if image_url:
                print(f"✅ 生成成功：{image_url}")
                return image_url
            else:
                print(f"❌ API响应中未找到图片URL")
                return None
        else:
            print(f"❌ API响应格式不正确或为空")
            return None

    except requests.exceptions.HTTPError as err:
        print(f"❌ HTTP请求异常（描述：{prompt[:30]}...）：{err.response.text}")
        return None
    except Exception as e:
        print(f"❌ 生成异常（描述：{prompt[:30]}...）：{str(e)}")
        return None


# -------------------------- 4. 批量生成函数（返回“描述-地址”关联列表） --------------------------
def batch_generate_images(prompts: List[str]) -> List[Dict[str, str]]:
    """
    返回格式：
    [
        {"prompt": "图片描述1", "image_path": "图片1URL"},
        {"prompt": "图片描述2", "image_path": "图片2URL"}
    ]
    仅包含生成成功的记录
    """
    # 核心：存储“描述-地址”关联的列表
    generated_result: List[Dict[str, str]] = []
    total = len(prompts)
    success_count = 0

    print(f"\n🚀 批量任务启动（共{total}张图片）")
    for idx, prompt in enumerate(prompts, 1):
        # 生成图片并关联描述
        image_url = generate_single_image(prompt)
        if image_url:
            # 新增“描述-地址”字典到列表
            generated_result.append({
                "prompt": prompt,
                "image_path": image_url
            })
            success_count += 1

    # 打印总结（展示关联关系）
    print(f"\n" + "=" * 60)
    print(f"📊 批量任务完成！")
    print(f"总计：{total}条描述 | 成功：{success_count}条 | 失败：{total - success_count}条")
    print(f"\n📋 成功生成的“描述-地址”列表：")
    for i, item in enumerate(generated_result, 1):
        print(f"\n{i}. 描述：{item['prompt']}")
        print(f"   地址：{item['image_path']}")
    print("=" * 60)

    return generated_result


# -------------------------- 5. 主执行入口（示例如何使用关联列表） --------------------------
if __name__ == "__main__":
    prompt_path_list = batch_generate_images(IMAGE_PROMPTS)

    # 示例1：遍历列表，根据描述筛选图片（如筛选“柯基”相关图片）
    if prompt_path_list:
        print(f"\n💡 示例1：筛选含“柯基”的图片")
        corgi_images = [item for item in prompt_path_list if "柯基" in item["prompt"]]
        if corgi_images:
            for item in corgi_images:
                print(f" - 描述：{item['prompt']} | 地址：{item['image_path']}")
        else:
            print(f" - 未找到含“柯基”的图片")

    # 示例2：遍历列表，批量处理图片（如批量下载、上传等）
    if prompt_path_list:
        print(f"\n💡 示例2：批量处理所有成功图片")
        for idx, item in enumerate(prompt_path_list, 1):
            print(f"\n处理第{idx}张：")
            print(f"  描述：{item['prompt']}")
            print(f"  图片URL：{item['image_path']}")
            # 这里可添加实际处理逻辑（如下载图片到本地）
            # response = requests.get(item['image_path'])
            # with open(f"image_{idx}.jpg", "wb") as f:
            #     f.write(response.content)
            # print(f"  已下载到本地：image_{idx}.jpg")
    else:
        print(f"\n💡 无成功生成的图片，无需后续处理")